package org.example.core.Service.impl;

import cn.hutool.core.date.DateTime;
import com.alibaba.fastjson2.JSON;
import org.example.core.Service.*;
import org.example.core.pojo.AllMessgae;
import org.example.core.pojo.ComfyuiTask;
import org.example.core.pojo.MessageBase;
import org.redisson.api.RSemaphore;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.example.core.job.RunTaskJob.TASK_RUN_SEMAPHORE;

@Service
public class ComfyuiMessageServiceImpl implements ComfyuiMessageService {
    @Autowired
    wsNoticeService wsNoticeService;
    @Autowired
    RedisService redisService;
    @Autowired
    UserResultService userResultService;
    @Autowired
    RedissonClient redissonClient;
    @Autowired
    UserFundRecordService userFundRecordService;
    @Override
    public void handleMessage(MessageBase messageBase) {
        if("progress".equals(messageBase.getType())){
            handlerProgressMessgae(messageBase);
        }
        if("executed".equals(messageBase.getType())){
            handlerExecutedMessgae(messageBase);
        }
        if("execution_error".equals(messageBase.getType())){
            handlerExecutionErrorMessgae(messageBase);
        }if("status".equals(messageBase.getType())){
            handlerStatusMessgae(messageBase);
        }
    }

    private void handlerStatusMessgae(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData();
        HashMap<String, Object> status = (HashMap<String, Object>)data.get("status");
        HashMap<String, Object> execInfo = (HashMap<String, Object>) status.get("exec_info");
        int i = Integer.valueOf(execInfo.get("queue_remaining") + "");
        if(i==0){
            //表示comfyui队列为空，可以添加任务
            RSemaphore semaphore = redissonClient.getSemaphore(TASK_RUN_SEMAPHORE);
            semaphore.release();//归还信号量
            System.out.println(DateTime.now()+"\t\t释放许可数量"+semaphore.availablePermits());
        }
    }

    private void handlerExecutionErrorMessgae(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData();
        data.put("type","execution_error");
        String promptId = (String) data.get("prompt_id");
        ComfyuiTask startTask = redisService.getStartTask(promptId);
        if (startTask == null){
            return;
        }
        userFundRecordService.freezeReturn(startTask.getUserId(),startTask.getSize());
        wsNoticeService.sendToUser(startTask.getWsClientId(), JSON.toJSONString(data));
        RSemaphore semaphore = redissonClient.getSemaphore(TASK_RUN_SEMAPHORE);
        semaphore.release();
    }

    private void handlerExecutedMessgae(MessageBase messageBase) {
        HashMap<String,Object> data =  messageBase.getData();
        HashMap<String,Object> output = (HashMap<String, Object>) data.get("output");
        List<Map<String,Object>> list = (List<Map<String, Object>>) output.get("images");
        List<String> urls = list.stream().map((image) ->
                        String.format("http://192.168.100.129:8188/view?filename=%s&type=%s&&subfolder", image.get("filename"), image.get("type"), image.get("subfolder")))
                .collect(Collectors.toList());

        HashMap<String,Object> result = new HashMap<>();
        result.put("type","imageResult");
        result.put("urls",urls);
        String promptId = (String) data.get("prompt_id");
        ComfyuiTask startTask = redisService.getStartTask(promptId);
        if (startTask == null){
            return;
        }
        userResultService.saveList(urls,startTask.getUserId());
        wsNoticeService.sendToUser(startTask.getWsClientId(), JSON.toJSONString(result));
        redisService.rmStartTask(promptId);
        RSemaphore semaphore = redissonClient.getSemaphore(TASK_RUN_SEMAPHORE);
        semaphore.release();
    }

    private void handlerProgressMessgae(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData();

        String promptId = (String) data.get("prompt_id");
        ComfyuiTask startTask = redisService.getStartTask(promptId);
        data.put("type","progress");
        if (startTask == null){
            return;
        }

        wsNoticeService.sendToUser(startTask.getWsClientId(), JSON.toJSONString(data));
        AllMessgae  allMessgae = new AllMessgae();
        allMessgae.setType("index");
        allMessgae.setValue(startTask.getIndex());
        wsNoticeService.sendToAll(JSON.toJSONString(allMessgae));
    }
}
